#没做过nlp项目，仅根据常规流程补充代码

import torch
import torch.nn as nn
from torch.nn import functional as F
import sys
sys.setrecursionlimit(3000)
torch.manual_seed(42)

vocab = list("YaH1%XfqFwLN(jZ'J,*OP_;.x8B-uI:GE|W=cAnDi&T#7y}K2gCQmUo56!M\9]l>k0`V^R[p+<$tr@{4S/b)dhsvz3e?")


class BiModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):

        logits = self.embedding(idx)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B * T, C)
            targets = targets.view(B * T)
            loss = F.cross_entropy(logits, targets)
        return logits

start = 'a'

def preding(word):

    idx = torch.tensor([[vocab.index(word)]], dtype=torch.long)
    logits = model(idx)
    outidx = logits.argmax(dim=-1).item()
    pred = vocab[outidx]
    if len(pred)<42:
        return preding(pred)
    else:
        print(pred)
        return pred


model = BiModel(len(vocab))
model.load_state_dict(torch.load('question3.pt',map_location=torch.device('cpu')))
with torch.no_grad():
    preding(start)
